// killproc.cpp
//
// Termination of processes.
//
// $Id: $
//

#include "stdafx.h"
#include "killproc.h"
#include <tlhelp32.h>

//---------------------------------------------------------------------------
// KillProcess
//
//  Terminates the specified process.
//
//  Parameters:
//	  dwProcessId - identifier of the process to terminate
//
//  Returns:
//	  TRUE, if successful, FALSE - otherwise.
//
BOOL
WINAPI
KillProcess(
	IN DWORD dwProcessId
	)
{
	HANDLE hProcess;
	DWORD dwError;

	// first try to obtain handle to the process without the use of any
	// additional privileges
	hProcess = OpenProcess(PROCESS_TERMINATE, FALSE, dwProcessId);
	if (hProcess == NULL)
	{
		if (GetLastError() != ERROR_ACCESS_DENIED)
			return FALSE;

		OSVERSIONINFO osvi;

		// determine operating system version
		osvi.dwOSVersionInfoSize = sizeof(osvi);
		GetVersionEx(&osvi);

		// we cannot do anything else if this is not Windows NT
		if (osvi.dwPlatformId != VER_PLATFORM_WIN32_NT)
			return SetLastError(ERROR_ACCESS_DENIED), FALSE;

		// enable SE_DEBUG_NAME privilege and try again

		TOKEN_PRIVILEGES Priv, PrivOld;
		DWORD cbPriv = sizeof(PrivOld);
		HANDLE hToken;

		// obtain the token of the current thread 
		if (!OpenThreadToken(GetCurrentThread(), 
							 TOKEN_QUERY|TOKEN_ADJUST_PRIVILEGES,
							 FALSE, &hToken))
		{
			if (GetLastError() != ERROR_NO_TOKEN)
				return FALSE;

			// revert to the process token
			if (!OpenProcessToken(GetCurrentProcess(),
								  TOKEN_QUERY|TOKEN_ADJUST_PRIVILEGES,
								  &hToken))
				return FALSE;
		}

		_ASSERTE(ANYSIZE_ARRAY > 0);

		Priv.PrivilegeCount = 1;
		Priv.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
		LookupPrivilegeValue(NULL, SE_DEBUG_NAME, &Priv.Privileges[0].Luid);

		// try to enable the privilege
		if (!AdjustTokenPrivileges(hToken, FALSE, &Priv, sizeof(Priv),
								   &PrivOld, &cbPriv))
		{
			dwError = GetLastError();
			CloseHandle(hToken);
			return SetLastError(dwError), FALSE;
		}

		if (GetLastError() == ERROR_NOT_ALL_ASSIGNED)
		{
			// the SE_DEBUG_NAME privilege is not present in the caller's
			// token
			CloseHandle(hToken);
			return SetLastError(ERROR_ACCESS_DENIED), FALSE;
		}

		// try to open process handle again
		hProcess = OpenProcess(PROCESS_TERMINATE, FALSE, dwProcessId);
		dwError = GetLastError();
		
		// restore the original state of the privilege
		AdjustTokenPrivileges(hToken, FALSE, &PrivOld, sizeof(PrivOld),
							  NULL, NULL);
		CloseHandle(hToken);

		if (hProcess == NULL)
			return SetLastError(FALSE), NULL;
	}

	// terminate the process
	if (!TerminateProcess(hProcess, (UINT)-1))
	{
		dwError = GetLastError();
		CloseHandle(hProcess);
		return SetLastError(dwError), FALSE;
	}

	CloseHandle(hProcess);

	// completed successfully
	return TRUE;
}

typedef LONG	NTSTATUS;
typedef LONG	KPRIORITY;

#define NT_SUCCESS(Status) ((NTSTATUS)(Status) >= 0)

#define STATUS_INFO_LENGTH_MISMATCH      ((NTSTATUS)0xC0000004L)

#define SystemProcessesAndThreadsInformation	5

typedef struct _CLIENT_ID {
    DWORD	    UniqueProcess;
    DWORD	    UniqueThread;
} CLIENT_ID;

typedef struct _UNICODE_STRING {
    USHORT	    Length;
    USHORT	    MaximumLength;
    PWSTR	    Buffer;
} UNICODE_STRING;

typedef struct _VM_COUNTERS {
    SIZE_T	    PeakVirtualSize;
    SIZE_T	    VirtualSize;
    ULONG	    PageFaultCount;
    SIZE_T	    PeakWorkingSetSize;
    SIZE_T	    WorkingSetSize;
    SIZE_T	    QuotaPeakPagedPoolUsage;
    SIZE_T	    QuotaPagedPoolUsage;
    SIZE_T	    QuotaPeakNonPagedPoolUsage;
    SIZE_T	    QuotaNonPagedPoolUsage;
    SIZE_T	    PagefileUsage;
    SIZE_T	    PeakPagefileUsage;
} VM_COUNTERS;

typedef struct _SYSTEM_THREADS {
    LARGE_INTEGER   KernelTime;
    LARGE_INTEGER   UserTime;
    LARGE_INTEGER   CreateTime;
    ULONG			WaitTime;
    PVOID			StartAddress;
    CLIENT_ID	    ClientId;
    KPRIORITY	    Priority;
    KPRIORITY	    BasePriority;
    ULONG			ContextSwitchCount;
    LONG			State;
    LONG			WaitReason;
} SYSTEM_THREADS, * PSYSTEM_THREADS;

// Note that the size of the SYSTEM_PROCESSES structure is different on
// NT 4 and Win2K, but we don't care about it, since we don't access neither
// IoCounters member nor Threads array

typedef struct _SYSTEM_PROCESSES {
    ULONG			NextEntryDelta;
    ULONG			ThreadCount;
    ULONG			Reserved1[6];
    LARGE_INTEGER   CreateTime;
    LARGE_INTEGER   UserTime;
    LARGE_INTEGER   KernelTime;
    UNICODE_STRING  ProcessName;
    KPRIORITY	    BasePriority;
    ULONG			ProcessId;
    ULONG			InheritedFromProcessId;
    ULONG			HandleCount;
    ULONG			Reserved2[2];
    VM_COUNTERS	    VmCounters;
#if _WIN32_WINNT >= 0x500
    IO_COUNTERS	    IoCounters;
#endif
    SYSTEM_THREADS  Threads[1];
} SYSTEM_PROCESSES, * PSYSTEM_PROCESSES;

//---------------------------------------------------------------------------
// KillProcessTreeNtHelper
//
//  This is a recursive helper function that terminates all the processes
//  started by the specified process and them terminates the process itself
//
//  Parameters:
//	  pInfo       - processes information
//	  dwProcessId - identifier of the process to terminate
//
//  Returns:
//	  Win32 error code.
//
static
BOOL
WINAPI
KillProcessTreeNtHelper(
	IN PSYSTEM_PROCESSES pInfo,
	IN DWORD dwProcessId
	)
{
	_ASSERTE(pInfo != NULL);

    PSYSTEM_PROCESSES p = pInfo;

    // kill all children first
    for (;;)
    {
		if (p->InheritedFromProcessId == dwProcessId)
			KillProcessTreeNtHelper(pInfo, p->ProcessId);

		if (p->NextEntryDelta == 0)
			break;

		// find the address of the next process structure
		p = (PSYSTEM_PROCESSES)(((LPBYTE)p) + p->NextEntryDelta);
    }

	// kill the process itself
    if (!KillProcess(dwProcessId))
		return GetLastError();

	return ERROR_SUCCESS;
}

//---------------------------------------------------------------------------
// KillProcessTreeWinHelper
//
//  This is a recursive helper function that terminates all the processes
//  started by the specified process and them terminates the process itself
//
//  Parameters:
//	  dwProcessId - identifier of the process to terminate
//
//  Returns:
//	  Win32 error code.
//
BOOL
WINAPI
KillProcessTreeWinHelper(
	IN DWORD dwProcessId
	)
{
	HINSTANCE hKernel;
	HANDLE (WINAPI * _CreateToolhelp32Snapshot)(DWORD, DWORD);
	BOOL (WINAPI * _Process32First)(HANDLE, PROCESSENTRY32 *);
	BOOL (WINAPI * _Process32Next)(HANDLE, PROCESSENTRY32 *);

	// get handle to KERNEL32.DLL
	hKernel = GetModuleHandle(_T("kernel32.dll"));
	_ASSERTE(hKernel != NULL);

	// locate necessary functions in KERNEL32.DLL
	*(FARPROC *)&_CreateToolhelp32Snapshot =
		GetProcAddress(hKernel, "CreateToolhelp32Snapshot");
	*(FARPROC *)&_Process32First =
		GetProcAddress(hKernel, "Process32First");
	*(FARPROC *)&_Process32Next =
		GetProcAddress(hKernel, "Process32Next");

	if (_CreateToolhelp32Snapshot == NULL ||
		_Process32First == NULL ||
		_Process32Next == NULL)
		return ERROR_PROC_NOT_FOUND;

	HANDLE hSnapshot;
	PROCESSENTRY32 Entry;

	// create a snapshot
	hSnapshot = _CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
	if (hSnapshot == INVALID_HANDLE_VALUE)
		return GetLastError();

	Entry.dwSize = sizeof(Entry);
	if (!_Process32First(hSnapshot, &Entry))
	{
		DWORD dwError = GetLastError();
		CloseHandle(hSnapshot);
		return dwError;
	}

	// kill all children first
	do
	{
		if (Entry.th32ParentProcessID == dwProcessId)
			KillProcessTreeWinHelper(Entry.th32ProcessID);

		Entry.dwSize = sizeof(Entry);

		TRACE("%d %d\n", Entry.th32ProcessID, Entry.th32ParentProcessID);
	}
	while (_Process32Next(hSnapshot, &Entry));

	CloseHandle(hSnapshot);

	// kill the process itself
    if (!KillProcess(dwProcessId))
		return GetLastError();

	return ERROR_SUCCESS;
}
BOOL
WINAPI
KillChildProcessTreeWinHelper(
	IN DWORD dwProcessId
	)
{
	HINSTANCE hKernel;
	HANDLE (WINAPI * _CreateToolhelp32Snapshot)(DWORD, DWORD);
	BOOL (WINAPI * _Process32First)(HANDLE, PROCESSENTRY32 *);
	BOOL (WINAPI * _Process32Next)(HANDLE, PROCESSENTRY32 *);

	// get handle to KERNEL32.DLL
	hKernel = GetModuleHandle(_T("kernel32.dll"));
	_ASSERTE(hKernel != NULL);

	// locate necessary functions in KERNEL32.DLL
	*(FARPROC *)&_CreateToolhelp32Snapshot =
		GetProcAddress(hKernel, "CreateToolhelp32Snapshot");
	*(FARPROC *)&_Process32First =
		GetProcAddress(hKernel, "Process32First");
	*(FARPROC *)&_Process32Next =
		GetProcAddress(hKernel, "Process32Next");

	if (_CreateToolhelp32Snapshot == NULL ||
		_Process32First == NULL ||
		_Process32Next == NULL)
		return ERROR_PROC_NOT_FOUND;

	HANDLE hSnapshot;
	PROCESSENTRY32 Entry;

	// create a snapshot
	hSnapshot = _CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
	if (hSnapshot == INVALID_HANDLE_VALUE)
		return GetLastError();

	Entry.dwSize = sizeof(Entry);
	if (!_Process32First(hSnapshot, &Entry))
	{
		DWORD dwError = GetLastError();
		CloseHandle(hSnapshot);
		return dwError;
	}

	// kill all children first
	do
	{
		if (Entry.th32ParentProcessID == dwProcessId)
			KillProcessTreeWinHelper(Entry.th32ProcessID);

		Entry.dwSize = sizeof(Entry);

		TRACE("%d %d\n", Entry.th32ProcessID, Entry.th32ParentProcessID);
	}
	while (_Process32Next(hSnapshot, &Entry));

	CloseHandle(hSnapshot);

	return ERROR_SUCCESS;
}
BOOL
WINAPI
KillParentProcessWinHelper(
	IN DWORD dwProcessId
	)
{
	HINSTANCE hKernel;
	HANDLE (WINAPI * _CreateToolhelp32Snapshot)(DWORD, DWORD);
	BOOL (WINAPI * _Process32First)(HANDLE, PROCESSENTRY32 *);
	BOOL (WINAPI * _Process32Next)(HANDLE, PROCESSENTRY32 *);

	// get handle to KERNEL32.DLL
	hKernel = GetModuleHandle(_T("kernel32.dll"));
	_ASSERTE(hKernel != NULL);

	// locate necessary functions in KERNEL32.DLL
	*(FARPROC *)&_CreateToolhelp32Snapshot =
		GetProcAddress(hKernel, "CreateToolhelp32Snapshot");
	*(FARPROC *)&_Process32First =
		GetProcAddress(hKernel, "Process32First");
	*(FARPROC *)&_Process32Next =
		GetProcAddress(hKernel, "Process32Next");

	if (_CreateToolhelp32Snapshot == NULL ||
		_Process32First == NULL ||
		_Process32Next == NULL)
		return ERROR_PROC_NOT_FOUND;

	HANDLE hSnapshot;
	PROCESSENTRY32 Entry;

	// create a snapshot
	hSnapshot = _CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
	if (hSnapshot == INVALID_HANDLE_VALUE)
		return GetLastError();

	Entry.dwSize = sizeof(Entry);
	if (!_Process32First(hSnapshot, &Entry))
	{
		DWORD dwError = GetLastError();
		CloseHandle(hSnapshot);
		return dwError;
	}

	//findParent ProcessID
	do
	{
		if (Entry.th32ProcessID == dwProcessId)
		{
			if (!KillProcess(Entry.th32ParentProcessID))
			{	return GetLastError();}
		}
		Entry.dwSize = sizeof(Entry);

		TRACE("%d %d\n", Entry.th32ProcessID, Entry.th32ParentProcessID);
	}
	while (_Process32Next(hSnapshot, &Entry));

	CloseHandle(hSnapshot);

	return ERROR_SUCCESS;
}

//---------------------------------------------------------------------------
// KillProcessEx
//
//  Terminates the specified process and, optionally, all processes started
//	from the specified process (the so-called process tree).
//
//  Parameters:
//	  dwProcessId - identifier of the process to terminate
//	  bTree		  - specifies whether the entire process tree should be
//					terminated
//
//  Returns:
//	  TRUE, if successful, FALSE - otherwise.
//
BOOL
WINAPI
KillProcessEx(
	IN DWORD dwProcessId,
	IN BOOL bTree
	)
{
	if (!bTree)
		return KillProcess(dwProcessId);

	OSVERSIONINFO osvi;
	DWORD dwError;

	// determine operating system version
	osvi.dwOSVersionInfoSize = sizeof(osvi);
	GetVersionEx(&osvi);

	if (osvi.dwPlatformId == VER_PLATFORM_WIN32_NT &&
		osvi.dwMajorVersion < 5)
	{
		HINSTANCE hNtDll;
		NTSTATUS (WINAPI * _ZwQuerySystemInformation)(UINT, PVOID, ULONG, PULONG);

		// get handle to NTDLL.DLL
		hNtDll = GetModuleHandle(_T("ntdll.dll"));
		_ASSERTE(hNtDll != NULL);

		// find the address of ZwQuerySystemInformation
		*(FARPROC *)&_ZwQuerySystemInformation =
			GetProcAddress(hNtDll, "ZwQuerySystemInformation");
		if (_ZwQuerySystemInformation == NULL)
			return SetLastError(ERROR_PROC_NOT_FOUND), NULL;

		// obtain a handle to the default process heap
		HANDLE hHeap = GetProcessHeap();
    
		NTSTATUS Status;
		ULONG cbBuffer = 0x8000;
		PVOID pBuffer = NULL;

		// it is difficult to say a priory which size of the buffer 
		// will be enough to retrieve all information, so we start
		// with 32K buffer and increase its size until we get the
		// information successfully
		do
		{
			pBuffer = HeapAlloc(hHeap, 0, cbBuffer);
			if (pBuffer == NULL)
				return SetLastError(ERROR_NOT_ENOUGH_MEMORY), FALSE;

			Status = _ZwQuerySystemInformation(
							SystemProcessesAndThreadsInformation,
							pBuffer, cbBuffer, NULL);

			if (Status == STATUS_INFO_LENGTH_MISMATCH)
			{
				HeapFree(hHeap, 0, pBuffer);
				cbBuffer *= 2;
			}
			else if (!NT_SUCCESS(Status))
			{
				HeapFree(hHeap, 0, pBuffer);
				return SetLastError(Status), NULL;
			}
		}
		while (Status == STATUS_INFO_LENGTH_MISMATCH);

		// call the helper function
		dwError = KillProcessTreeNtHelper((PSYSTEM_PROCESSES)pBuffer, 
										  dwProcessId);
		
		HeapFree(hHeap, 0, pBuffer);
	}
	else
	{
		// call the helper function
		dwError = KillProcessTreeWinHelper(dwProcessId);
	}

	SetLastError(dwError);
	return dwError == ERROR_SUCCESS;
}

DWORD ProcessKill(CString strProcessName)

{
    HANDLE         hProcessSnap = NULL;
    DWORD          dwReturn     = 0;
    PROCESSENTRY32 pe32         = {0};

    strProcessName.MakeLower();
    hProcessSnap = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);


    if(hProcessSnap == INVALID_HANDLE_VALUE)
        return (DWORD)INVALID_HANDLE_VALUE;

    pe32.dwSize = sizeof(PROCESSENTRY32);

    if(Process32First(hProcessSnap, &pe32))
    {
        DWORD Code = 0;
        DWORD dwPriorityClass;

        do
        {
            HANDLE hProcess;
            hProcess = OpenProcess(PROCESS_ALL_ACCESS, FALSE, pe32.th32ProcessID);
            dwPriorityClass = GetPriorityClass(hProcess);

            CString Temp = pe32.szExeFile;
            Temp.MakeLower();

            if(Temp == strProcessName)
            {
                if(TerminateProcess(hProcess, 0))
                    GetExitCodeProcess(hProcess, &Code);
                else
                    return GetLastError();
            }
            CloseHandle(hProcess);
        } while(Process32Next(hProcessSnap, &pe32));
    }
    CloseHandle(hProcessSnap); 
    return dwReturn;
}
